# %%
import json
from collections import defaultdict
from pydantic import BaseModel, Field
from typing import Literal
from datetime import datetime, timedelta
from connections import mongo_db, get_redis_conn

clip_col = mongo_db.clip
bag_col = mongo_db.bag
clip_annotated_record = mongo_db.clip_annotated_record
redis_client = get_redis_conn()


SCENE_TAGS = ["highway", "open_parking", "indoor_parking", "city_road"]
START_DATE = "2025-03-01"


class SceneClipCnt(BaseModel):
    collect_date: str = Field(..., description="采集日期")
    scene: Literal["highway", "city_road", "open_parking", "indoor_parking"] = Field(
        ..., description="线路场景标签"
    )
    collect_cnt: int = Field(default=0, description="采集量")
    parsed_cnt: int = Field(default=0, description="解析量")
    lidar_slam_cnt: int = Field(default=0, description="Lidar+SLAM解析量")
    ready_curated_cnt: int = Field(default=0, description="待送标量")
    curated_cnt: int = Field(default=0, description="已送标量")
    not_curated_cnt: int = Field(default=0, description="未送标量")


def summary_by_date(date: str):
    # 不同场景的采集量、解析量、Lidar+SLAM解析量、待送标量、已送标量、未送标量
    clip_cnt_by_scene = {
        tag: SceneClipCnt(collect_date=date, scene=tag) for tag in SCENE_TAGS
    }
    # 不同场景的待送标量 clip_id 列表，方便从 clip_annotated_record 中统计已送标量
    curated_clip_ids_by_scene = {tag: [] for tag in SCENE_TAGS}

    # 采集量，从 bag 中统计
    bags = bag_col.find(
        {"collect_at": {"$regex": rf"^{date}"}}, {"tags.route": 1, "collect_at": 1}
    ).to_list()
    for bag in bags:
        assert bag["collect_at"].startswith(date)
        route_tags = bag.get("tags", {}).get("route", [])
        for key_tag in SCENE_TAGS:
            if key_tag in route_tags:
                clip_cnt_by_scene[key_tag].collect_cnt += 1
                break

    # 解析量，做过 lidar slam 的量，可以送标量，都从 clip 中统计
    clips = clip_col.find(
        {"collect_at": {"$regex": rf"^{date}"}},
        {
            "tags.route": 1,
            "tags.self_posture": 1,
            "ext_info.lidar0": 1,
            "collect_at": 1,
        },
    ).to_list()
    for clip in clips:
        assert clip["collect_at"].startswith(date)
        route_tags = clip.get("tags", {}).get("route", [])
        self_posture = clip.get("tags", {}).get("self_posture", [])
        ext_info_lidar0 = clip.get("ext_info", {}).get("lidar0", False)

        # 一旦找到 tag，则 break，即一个 clip 只能属于一个 tag
        # parsed_cnt, lidar_slam_cnt, ready_curated_cnt 是递进关系，即 parsed_cnt >= lidar_slam_cnt >= ready_curated_cnt
        for key_tag in SCENE_TAGS:
            if key_tag in route_tags:
                clip_cnt_by_scene[key_tag].parsed_cnt += 1
                if ext_info_lidar0:
                    clip_cnt_by_scene[key_tag].lidar_slam_cnt += 1
                    if "stop" not in self_posture:
                        clip_cnt_by_scene[key_tag].ready_curated_cnt += 1
                        curated_clip_ids_by_scene[key_tag].append(str(clip["_id"]))
                break

    # 从 clip_annotated_record 中统计已送标量
    for key_tag in SCENE_TAGS:
        clip_ids = curated_clip_ids_by_scene[key_tag]
        if clip_ids:
            curated_clip_ids = clip_annotated_record.distinct(
                "clip_id",
                {"clip_id": {"$in": clip_ids}},
            )
            clip_cnt_by_scene[key_tag].curated_cnt = len(curated_clip_ids)
            clip_cnt_by_scene[key_tag].not_curated_cnt = len(clip_ids) - len(
                curated_clip_ids
            )

    return list(clip_cnt_by_scene.values())


def main():
    cur_date = datetime.strptime(START_DATE, "%Y-%m-%d")
    clip_cnt_list = []
    while True:
        cur_date_str = cur_date.strftime("%Y-%m-%d")
        clip_cnt_list.extend(summary_by_date(cur_date_str))
        print("handle date success", cur_date_str)
        cur_date += timedelta(days=1)
        if cur_date > datetime.now():
            break

    # 按 scene 分类
    clip_cnt_by_scene = defaultdict(list)
    for clip_cnt in clip_cnt_list:
        clip_cnt_by_scene[clip_cnt.scene].append(clip_cnt)

    # 排序并放入缓存
    for k, v in clip_cnt_by_scene.items():
        v.sort(key=lambda x: x.collect_date)
        redis_client.set(
            f"portal:data_loop_dashboard:scene_clip_statistics:{k}",
            json.dumps([i.model_dump() for i in v]),
        )


if __name__ == "__main__":
    main()
